package solution.design;

import common.ListNode;

import java.util.LinkedList;
import java.util.List;

/**
 * @author zhangmin
 * @create 2022-05-23 9:46
 *
 * 设计--Trie树实现
 */

/* Trie 树节点实现 */
class TrieNode<V> {
    V val = null; //val 字段存储键对应的值
    TrieNode<V>[] children = new TrieNode[256]; //children 数组存储指向子节点的指针。
}
// 底层用 Trie 树实现的键值映射
// 键为 String 类型，值为类型 V
class TrieMap<V> {
    // ASCII 码个数
    private static final int R = 256;
    // 当前存在 Map 中的键值对个数
    private int size = 0;
    // Trie 树的根节点
    private TrieNode<V> root = null;
    private static class TrieNode<V> {
        V val = null;
        TrieNode<V>[] children = new TrieNode[R];
    }

    /* 其他 API 的实现... */
    //获得树中的所有节点数量
    public int getSize() {
        return size;
    }
    //工具方法：从node节点开始搜索key，如果存在返回对应的节点，否则返回null
    //注意：就算getNode(key)的返回值x非空，也只能说字符串key是一个「前缀」；
    //除非x.val同时非空，才能判断键key存在--可以用来实现hasKeyWithPrefix
    private TrieNode<V> getNode(TrieNode<V> node, String key) {
        TrieNode<V> p = node;
        // 从节点 node 开始搜索 key
        for (int i = 0; i < key.length(); i++) {
            if (p == null) {
                // 无法向下搜索
                return null;
            }
            // 向下搜索
            char c = key.charAt(i);
            p = p.children[c];
        }
        return p;
    }

    /***** 增/改 *****/
    // 在 Map 中添加 key
    public void put(String key, V val){
        if (!containsKey(key)){
            size++;
        }
        root=put(root,key,val,0);
    }
    // 定义：向以 node 为根的 Trie 树中插入 key[i..]，返回插入完成后的根节点
    private TrieNode<V> put(TrieNode<V> node,String key,V val,int i){
        if (node==null){
            //如果树枝不存在，就新建
            node=new TrieNode<>();
        }
        if (i==key.length()){
            //key 的路径已插入完成，将值 val 存入节点
            node.val=val;
            return node;
        }
        char c=key.charAt(i);
        // 递归插入子节点，并接收返回值
        node.children[c]=put(node.children[c],key,val,i+1);
        return node;
    }

    /***** 删 *****/
    // 删除键 key 以及对应的值
    //首先，递归修改数据结构的时候，如果一个节点想删掉自己，直接返回空指针就行了。
    //其次，一个节点如何知道自己是否需要被删除呢？
    //主要看自己的val字段是否为空以及自己的children数组是否全都是空指针。
    public void remove(String key){
        if (!containsKey(key)){
            return;
        }
        // 递归修改数据结构要接收函数的返回值
        root=remove(root,key,0);
        size--;
    }
    private TrieNode<V> remove(TrieNode<V> node,String key,int i){
        if (node==null) return null;
        if (i==key.length()){
            // 找到了 key 对应的 TrieNode，删除 val
            node.val=null;
        }else {
            char c=key.charAt(i);
            node.children[c]=remove(node.children[c],key,i+1);
        }
        //后序位置，递归路径上的节点可能需要被清理
        if (node.val!=null){
            // 如果该 TireNode 存储着 val，不需要被清理
            return node;
        }
        // 检查该 TrieNode 是否还有后缀
        for (int c=0;c<R;c++){
            if (node.children[c]!=null){
                // 只要存在一个子节点（后缀树枝），就不需要被清理
                return node;
            }
        }
        // 既没有存储 val，也没有后缀树枝，则该节点需要被清理
        return null;
    }

    /***** 查 *****/
    // 搜索 key 对应的值，不存在则返回 null
    // get("the") -> 4
    // get("tha") -> null
    public V get(String key){
        // 从 root 开始搜索 key
        TrieNode<V> x = getNode(root, key);
        if (x == null || x.val == null) {
            // x 为空或 x 的 val 字段为空都说明 key 没有对应的值
            return null;
        }
        return x.val;
    }

    // 判断 key 是否存在在 Map 中
    // containsKey("tea") -> false
    // containsKey("team") -> true
    public boolean containsKey(String key){
        return get(key)!=null;
    }

    // 在 Map 的所有键中搜索 query 的最短前缀
    // shortestPrefixOf("themxyz") -> "the"
    public String shortestPrefixOf(String query){
        TrieNode<V> p=root;
        for (int i = 0; i < query.length(); i++) {
            if (p==null){
                //无法向下搜索
                return "";
            }
            if (p.val!=null){
                // 找到一个键是 query 的前缀
                return query.substring(0,i);
            }
            //继续向下搜索
            char c= query.charAt(i);
            p=p.children[c];
        }
        //需要检查最后一个节点
        if (p!=null&&p.val!=null){
            // 如果 query 本身就是一个键
            return query;
        }
        return "";
    }

    // 在 Map 的所有键中搜索 query 的最长前缀
    // longestPrefixOf("themxyz") -> "them"
    public String longestPrefixOf(String query){
        TrieNode<V> p=root;
        int max_len=0;
        for (int i = 0; i < query.length(); i++) {
            if (p==null){
                break;
            }
            if (p.val!=null){
                // 找到一个键是 query 的前缀，更新前缀的最大长度
                max_len=i;
            }
            char c=query.charAt(i);
            p=p.children[c];
        }
        if (p!=null&&p.val!=null){
            return query;
        }
        return query.substring(0,max_len);
    }

    // 搜索所有前缀为 prefix 的键
    // keysWithPrefix("th") -> ["that", "the", "them"]
    //先利用getNode函数在 Trie 树中找到prefix对应的节点x，然后开始多叉树的遍历算法，
    //遍历以x为根的这棵 Trie 树，找到所有键值对
    public List<String> keysWithPrefix(String prefix){
        List<String > res=new LinkedList<>();
        // 找到匹配 prefix 在 Trie 树中的那个节点
        TrieNode<V> x = getNode(root, prefix);
        if (x==null){
            return res;
        }
        // DFS 遍历以 x 为根的这棵 Trie 树
        traverse(x,new StringBuffer(prefix),res);
        return res;
    }
    // 遍历以 node 节点为根的 Trie 树，找到所有键
    private void traverse(TrieNode<V> node,StringBuffer path,List<String> res){
        if (node==null){
            return;
        }
        if (node.val!=null){
            res.add(path.toString());
        }
        // 回溯算法遍历框架
        for (char c = 0;  c< R; c++) {
            path.append(c);
            traverse(node.children[c],path,res);
            path.deleteCharAt(path.length()-1);
        }
    }

    // 判断是否存在前缀为 prefix 的键
    // hasKeyWithPrefix("tha") -> true
    // hasKeyWithPrefix("apple") -> false
    public boolean hasKeyWithPrefix(String prefix){
        return getNode(root,prefix)!=null;
    }

    // 通配符 . 匹配任意字符，搜索所有匹配的键
    // keysWithPattern("t.a.") -> ["team", "that"]
    public List<String> keysWithPattern(String pattern){
        List<String> res=new LinkedList<>();
        traverse(root,new StringBuffer(),pattern,0,res);
        return res;
    }
    // 遍历函数，尝试在「以 node 为根的 Trie 树中」匹配 pattern[i..]
    private void traverse(TrieNode<V> node,StringBuffer path,String pattern,int i,List<String> res){
        if (node==null){
            return;
        }
        if (i==pattern.length()){
            // pattern 匹配完成
            if (node.val!=null){
                // 如果这个节点存储着 val，则找到一个匹配的键
                res.add(path.toString());
            }
            return;
        }
        char c=pattern.charAt(i);
        if (c=='.'){
            // pattern[i] 是通配符，可以变化成任意字符
            // 多叉树（回溯算法）遍历框架
            for (int j = 0; j < R; j++) {
                path.append(j);
                traverse(node.children[j],path,pattern,i+1,res);
                path.deleteCharAt(path.length()-1);
            }
        }else {
            path.append(c);
            traverse(node.children[c],path,pattern,i+1,res);
            path.deleteCharAt(path.length()-1);
        }
    }

    // 通配符 . 匹配任意字符，判断是否存在匹配的键
    // hasKeyWithPattern(".ip") -> true
    // hasKeyWithPattern(".i") -> false
    public boolean hasKeyWithPattern(String pattern){
        return hasKeyWithPattern(root,pattern,0);
    }
    // 函数定义：从 node 节点开始匹配 pattern[i..]，返回是否成功匹配
    private boolean hasKeyWithPattern(TrieNode<V> node,String pattern,int i){
        if (node==null){
            // 树枝不存在，即匹配失败
            return false;
        }
        if (i==pattern.length()){
            // 模式串走到头了，看看匹配到的是否是一个键
            return node.val!=null;
        }
        char c=pattern.charAt(i);
        // 没有遇到通配符
        if (c!='.'){
            return hasKeyWithPattern(node.children[c],pattern,i+1);
        }
        // 遇到通配符
        for (int j = 0; j < R; j++) {
            // pattern[i] 可以变化成任意字符，尝试所有可能，只要遇到一个匹配成功就返回
            if (hasKeyWithPattern(node.children[j],pattern,i+1)){
                return true;
            }
        }
        return false;
    }
}
